from super_gradients import setup_device
from super_gradients.training import Trainer
from super_gradients.training.datasets import YoloDarknetFormatDetectionDataset
from super_gradients.training.losses import YoloXDetectionLoss, PPYoloELoss
from super_gradients.training.metrics import DetectionMetrics_050
from super_gradients.training.models.detection_models.pp_yolo_e import PPYoloEPostPredictionCallback
from super_gradients.training.transforms.transforms import DetectionRescale
from torch.utils.data import DataLoader
from super_gradients.training.utils.detection_utils import DetectionCollateFN
from super_gradients.training import models
from prettyformatter import pprint

CHECKPOINT_DIR = '.\\test'


def train():
    trainer = Trainer(experiment_name='breast_detection',
                      ckpt_root_dir=CHECKPOINT_DIR)
    train_dataset = YoloDarknetFormatDetectionDataset(data_dir="datasets\\yolo-darknet",
                                                      images_dir="train2007", labels_dir="train2007",
                                                      classes=["belt_logo"],
                                                      transforms=[DetectionRescale(output_shape=(1280, 640))])
    val_dataset = YoloDarknetFormatDetectionDataset(data_dir="datasets\\yolo-darknet",
                                                    images_dir="val2007", labels_dir="val2007",
                                                    classes=["belt_logo"],
                                                    transforms=[DetectionRescale(output_shape=(1280, 640))])

    test_dataset = YoloDarknetFormatDetectionDataset(data_dir="datasets\\yolo-darknet",
                                                     images_dir="test2007", labels_dir="test2007",
                                                     classes=["belt_logo"],
                                                     transforms=[DetectionRescale(output_shape=(1280, 640))])

    train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=0,
                                  collate_fn=DetectionCollateFN())
    val_dataloader = DataLoader(val_dataset, batch_size=4, shuffle=False, num_workers=0,
                                collate_fn=DetectionCollateFN())
    test_dataloader = DataLoader(test_dataset, batch_size=4, shuffle=True, num_workers=0,
                                 collate_fn=DetectionCollateFN())

    model = models.get("yolo_nas_s", pretrained_weights="coco", num_classes=1)
    train_params = {
        # ENABLING SILENT MODE
        'silent_mode': False,
        "average_best_models": True,
        "warmup_mode": "linear_epoch_step",
        "warmup_initial_lr": 1e-6,
        "lr_warmup_epochs": 3,
        "initial_lr": 5e-4,
        "lr_mode": "cosine",
        "cosine_final_lr_ratio": 0.1,
        "optimizer": "Adam",
        "optimizer_params": {"weight_decay": 0.0001},
        "zero_weight_decay_on_bias_and_bn": True,
        "ema": True,
        "ema_params": {"decay": 0.9, "decay_type": "threshold"},
        # ONLY TRAINING FOR 10 EPOCHS FOR THIS EXAMPLE NOTEBOOK
        "max_epochs": 20,
        "mixed_precision": True,
        "loss": PPYoloELoss(
            use_static_assigner=False,
            # NOTE: num_classes needs to be defined here
            num_classes=1,
            reg_max=16
        ),
        "valid_metrics_list": [
            DetectionMetrics_050(
                score_thres=0.1,
                top_k_predictions=300,
                # NOTE: num_classes needs to be defined here
                num_cls=1,
                normalize_targets=True,
                post_prediction_callback=PPYoloEPostPredictionCallback(
                    score_threshold=0.01,
                    nms_top_k=1000,
                    max_predictions=300,
                    nms_threshold=0.7
                )
            )
        ],
        "metric_to_watch": "mAP@0.50:0.95"
    }
    # train_params["num_gpus"] = 1
    pprint(train_params, json=True)
    setup_device(num_gpus=1)
    trainer.train(model=model, training_params=train_params, train_loader=train_dataloader, valid_loader=val_dataloader,
                  test_loaders=dict(test_dataloader))


if __name__ == "__main__":
    train()
